# optimal_knn.py - 使用 Pinecone 向量数据库实现 KNN 分类

import pickle
import matplotlib.pyplot as plt
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from tqdm import tqdm  # 打印进度条
from pinecone import Pinecone, ServerlessSpec
import numpy as np

# 加载手写数字数据集
digits = load_digits()
X, y = digits.data, digits.target

# 划分训练集与测试集
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# Pinecone 初始化
pc = Pinecone(api_key="pcsk_5kxYS1_7Wqt2wKhRJntVmYMMVVEvSxbYVJCTEM3Z1RK3viMkFVpjFLTJyA4xaHZdhGQajC")
index_name = "optimal-knn-index"

# 检查并删除现有索引（如果存在）
if index_name in pc.list_indexes().names():
    pc.delete_index(index_name)

# 创建新索引
pc.create_index(
    name=index_name,
    dimension=64,  # MNIST 8x8 = 64 维
    metric="euclidean",
    spec=ServerlessSpec(
        cloud="aws",
        region="us-east-1"
    )
)

# 连接到索引
index = pc.Index(index_name)

# 上传训练数据到 Pinecone
print("📤 正在上传训练数据到 Pinecone...")
batch_size = 100
for i in tqdm(range(0, len(X_train), batch_size), desc="上传数据"):
    batch_end = min(i + batch_size, len(X_train))
    vectors = []
    for j in range(i, batch_end):
        vectors.append({
            "id": str(j),
            "values": X_train[j].tolist(),
            "metadata": {"label": int(y_train[j])}
        })
    index.upsert(vectors)

print("✅ 训练数据上传完成！")

# 初始化变量
best_accuracy = 0.0
best_k = 1
accuracies = []

# 使用 Pinecone 进行 KNN 分类测试
print("🔍 正在使用 Pinecone 进行 KNN 分类测试...")
print(f"📊 测试样本数量: {len(X_test)}, K值范围: 1-20")

# 减少测试范围以提高效率
test_k_values = range(1, 21)  # 只测试 K=1 到 K=20
sample_size = min(100, len(X_test))  # 使用最多100个样本进行测试
X_test_sample = X_test[:sample_size]
y_test_sample = y_test[:sample_size]

print(f"🎯 使用 {sample_size} 个测试样本进行快速评估")

for k in tqdm(test_k_values, desc="正在测试不同的K值"):
    y_pred = []
    
    # 对每个测试样本进行查询
    for i, test_vector in enumerate(X_test_sample):
        # 在 Pinecone 中查询 k 个最近邻
        results = index.query(
            vector=test_vector.tolist(),
            top_k=k,
            include_metadata=True
        )
        
        # 提取邻居标签
        neighbor_labels = []
        if results['matches']:
            for match in results['matches']:
                if 'metadata' in match and 'label' in match['metadata']:
                    neighbor_labels.append(match['metadata']['label'])
        
        # 投票决定预测结果
        if neighbor_labels:
            # 统计每个标签的出现次数
            label_counts = {}
            for label in neighbor_labels:
                label_counts[label] = label_counts.get(label, 0) + 1
            
            # 选择出现次数最多的标签
            pred_label = max(label_counts, key=label_counts.get)
        else:
            pred_label = 0  # 如果没有邻居，默认预测为0
        
        y_pred.append(pred_label)
    
    # 计算准确率
    acc = accuracy_score(y_test_sample, y_pred)
    accuracies.append(acc)
    
    if acc > best_accuracy:
        best_accuracy = acc
        best_k = k

# 保存最佳配置
best_config = {
    'best_k': best_k,
    'best_accuracy': best_accuracy,
    'index_name': index_name,
    'total_vectors': len(X_train)
}

with open("best_pinecone_config.pkl", "wb") as f:
    pickle.dump(best_config, f)

print(f"\n✅ 最佳K值: {best_k}, 对应准确率: {best_accuracy:.4f}")
print(f"💾 Pinecone 索引名称: {index_name}")
print(f"📊 训练向量总数: {len(X_train)}")
print(f"🧪 测试样本数量: {len(X_test)}")

# 绘图
plt.figure(figsize=(10, 6))
plt.plot(list(test_k_values), accuracies, marker='o', linestyle='-', label='Accuracy')
plt.axvline(x=best_k, color='red', linestyle='--', label=f'Best K = {best_k}')
plt.scatter(best_k, best_accuracy, color='red', zorder=5)
plt.text(best_k + 0.5, best_accuracy, f"K={best_k}\nAcc={best_accuracy:.4f}", color='red')
plt.title("Pinecone KNN Accuracy vs K Value")
plt.xlabel("K value")
plt.ylabel("Accuracy")
plt.legend()
plt.grid(True)
plt.tight_layout()

# 保存为 PDF 文件
plt.savefig("knn_results.pdf")
plt.close()

print("📊 已保存准确率折线图为 knn_results.pdf")
print("💾 已保存最佳配置为 best_pinecone_config.pkl")
print("🎯 Pinecone KNN 模型优化完成！")
